import time
from threading import Lock
from functools import wraps


class RateLimiter:
    def __init__(self, max_requests_per_interval, interval):
        self.max_requests_per_interval = max_requests_per_interval
        self.interval = interval
        self.lock = Lock()
        self.requests = []

    def _cleanup_old_requests(self):
        """清理旧的请求时间戳"""
        current_time = time.time()
        self.requests = [timestamp for timestamp in self.requests if current_time - timestamp < self.interval]

    def _wait_if_needed(self):
        """根据请求速率限制等待时间"""
        with self.lock:
            self._cleanup_old_requests()
            if len(self.requests) >= self.max_requests_per_interval:
                wait_time = self.interval - (time.time() - self.requests[0])
                if wait_time > 0:
                    time.sleep(wait_time)
            self.requests.append(time.time())

    def limit(self, max_requests_per_interval=None, interval=None):
        """带参数的限流装饰器"""
        if max_requests_per_interval is None:
            max_requests_per_interval = self.max_requests_per_interval
        if interval is None:
            interval = self.interval

        def decorator(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                self.max_requests_per_interval = max_requests_per_interval
                self.interval = interval
                self._wait_if_needed()
                return func(*args, **kwargs)

            return wrapper

        return decorator
